package datastructure.tree.redblack;

import static datastructure.tree.redblack.RedBlackTree.Node.Color.BLACK;
import static datastructure.tree.redblack.RedBlackTree.Node.Color.RED;

/**
 * 红黑树
 *
 * @author 杨 强
 * @package datastructure.tree.redblack
 * @createTime 2023-04-06 20:59
 */
public class RedBlackTree {
    public static void main(String[] args) {
        RedBlackTree tree = new RedBlackTree();
        tree.put(8, 8);
        tree.put(10, 10);
        tree.put(5, 5);
        tree.put(2, 2);
        tree.put(3, 3);
        System.out.println(true);
    }

    /**
     * 红黑树根节点
     */
    private Node root;

    /**
     * 新增或更新
     * @param key 键
     * @param value 值
     */
    public void put(int key, Object value){
        // 找到新增节点的父节点
        Node p = root;
        Node parent = null;
        while (p != null) {
            parent = p;
            if (key < p.key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                // key相等就更新退出
                p.value = value; // 更新
                return;
            }
        }
        // 创建新增节点
        Node inserted = new Node(key, value);
        if (parent == null) {
            // 没有父节点,那么新增节点做根节点
            root = inserted;
        } else if (key < parent.key) {
            // 新增节点在父节点左边
            parent.left = inserted;
            inserted.parent = parent;
        } else {
            // 新增节点在父节点右边
            parent.right = inserted;
            inserted.parent = parent;
        }
        // 新增进行调整平衡
        fixRedRed(inserted);
    }


    /**
     * 插入节点维护平衡的方法
     *     主要是修复【红红】的不平衡情况,然后红红相邻不平衡会出现很多情况
     * @param x 新插入的节点
     */
    private void fixRedRed(Node x){
        // case 1 插入节点是根节点, 变黑即可
        if (x == root){
            x.color = BLACK;
            return;
        }
        // case 2 插入节点父亲节点是黑色,无需调整
        if (isBlack(x.parent)){
            return;
        }
        /*
            case 3 当父节点为红色时,触发红红相邻,并且叔叔节点也为红色.
                   需要将父亲叔叔变黑,然祖父变红,然后对祖父进行递归
         */
        Node parent = x.parent;// 父
        Node uncle = x.uncle();// 叔
        Node grandParent = parent.parent;// 祖父
        if (isRed(uncle)){// 叔为红
            parent.color = BLACK;
            uncle.color = BLACK;
            grandParent.color = RED;
            fixRedRed(grandParent);
            return;
        }

         /*
            case 4 当父节点为红色时,触发红红相邻,并且叔叔节点为黑色节点
         */
        if (parent.isLeftChild() && x.isLeftChild()){
            // 父亲是左孩子,插入节点也是左孩子,就是情况4的第一种
            // 然后父亲变黑,祖父变红
            parent.color = BLACK;
            grandParent.color = RED;
            // 然后祖父右旋
            rightRotate(grandParent);
        }else if (parent.isLeftChild() && !x.isLeftChild()){
            // 父亲是左孩子,插入节点是右孩子,就是情况4的第二种
            // 父亲节点先左旋
            leftRotate(parent);
            // 然后变色
            x.color=BLACK;
            grandParent.color=RED;
            // 祖父右旋
            rightRotate(grandParent);
        }else if (!parent.isLeftChild() && !x.isLeftChild()){
            // 右右RR 不平衡
            parent.color = BLACK;
            grandParent.color = RED;
            leftRotate(grandParent);
        }else {
            // 右左RL 不平衡
            rightRotate(parent);
            x.color=BLACK;
            grandParent.color=RED;
            leftRotate(grandParent);
        }
    }





    // 判断节点是否是红色
    private boolean isRed(Node node){
        return node != null && node.color == RED;
    }

    // 判断节点是否是黑色
    private boolean isBlack(Node node){
        return node == null || node.color == BLACK;
    }

    /**
     * 右旋
     * 处理parent,处理新根节点的父子关系
     * @param pink 粉色节点:需要进行旋转的子树的根节点
     */
    private void rightRotate(Node pink){
        // 旋转后的根节点的父节点
        Node parent = pink.parent;

        // 进行旋转,并且处理旋转后的节点的parent节点
        Node yellow = pink.left;
        Node green = yellow.right;
        if (green != null){
            green.parent = pink;
        }
        yellow.right = pink;
        yellow.parent =parent;
        pink.left = green;
        pink.parent = yellow;

        // 处理新根节点的父节点
        if (parent == null){
            root = yellow;
        }else if (parent.left == pink){
            parent.left = yellow;
        }else {
            parent.right = yellow;
        }
    }

    /**
     * 左旋
     * 处理parent,处理新根节点的父子关系
     * @param pink 粉色节点:需要进行旋转的子树的根节点
     */
    private void leftRotate(Node pink){
        Node parent = pink.parent;
        Node yellow = pink.right;
        Node green = yellow.left;
        if (green != null) {
            green.parent = pink;
        }
        yellow.left = pink;
        yellow.parent = parent;
        pink.right = green;
        pink.parent = yellow;
        if (parent == null) {
            root = yellow;
        } else if (parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }


    /**
     * 节点类
     */
    public static class Node {
        /**
         * 颜色枚举
         */
        enum Color {
            RED, BLACK;
        }

        /**
         * key
         */
        int key;
        /**
         * 节点元素值
         */
        Object value;
        /**
         * 左节点
         */
        Node left;
        /**
         * 右节点
         */
        Node right;
        /**
         * 节点颜色
         */
        Color color = Color.RED;
        /**
         * 父节点
         */
        Node parent;

        public Node(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public Node(int key) {
            this.key = key;
        }

        public Node(int key, Color color) {
            this.key = key;
            this.color = color;
        }

        public Node(int key, Color color, Node left, Node right) {
            this.key = key;
            this.color = color;
            this.left = left;
            this.right = right;
            // 如果节点的左孩子节点不为空,那么左孩子节点的父节点就是当前节点
            if (left != null) {
                left.parent = this;
            }
            // 如果节点的右孩子节点不为空,那么右孩子节点的父节点就是当前节点
            if (right != null) {
                right.parent = this;
            }
        }

        // 是否是左孩子
        boolean isLeftChild(){
            return parent != null && parent.left == this;
        }

        // 找节点的叔叔节点
        Node uncle(){
            if (parent == null || parent.parent == null){
                return null;
            }
            // 如果节点的父节点是左节点,那么叔叔节点就是父节点的父节点的右节点
            if (parent.isLeftChild()){
                return parent.parent.right;
            }else {
                return parent.parent.left;
            }

        }

        // 找节点的兄弟节点
        Node sibling(){
            if (parent == null){
                return null;
            }
            if (this.isLeftChild()){
                return parent.right;
            }else {
                return parent.left;
            }
        }

    }

}
